import numpy as np
import matplotlib.pyplot as plt
def solFunc(n_bits):
    return [np.random.randint(0,2) for j in range(n_bits)]

#you should judge the solution if valid
def accept(sol,table, W):
    cost = 0
    index = 0
    for i in (sol):
        if i == 1:
            cost += table[i][0]
        index += 1
    if cost > W:
        return False
    else: return True
def objFunc(sol, table, W):
    value = 0
    index = 0
    for i in sol:
        if i == 1:
            value += table[index][1]
        index += 1
    if not accept(sol, table, W):
        value = -1
    return value


class GA(object):
    def __init__(self,solFunc,objFunc,table,W,n_bits, n_iter, n_pop, r_cross, r_mut):
        self.n_bits = n_bits
        self.n_iter = n_iter
        self.n_pop = n_pop
        self.r_cross = r_cross
        self.r_mut = r_mut
        #function as parameter
        self.objFunc = objFunc
        self.solFunc = solFunc
        #hanspack
        self.W = W
        self.table = table
        #init poplution, and get scores 
        self.pop = [self.solFunc(n_bits) for i in range(self.n_pop)]
        self.scores =[self.objFunc(c,self.table, self.W) for c in self.pop]#scores of popution
        #learning rate
        self.learnRate = list()
        #solution
        self.best = self.pop[0]
        self.best_eval =-1# self.objFunc(self.best, self.table, self.W)
        self.index = list()
    def getIndexObj(self):
        index = 0
        for i in self.best:
            if i == 1:
                self.index.append(index)
            index += 1

    def selection(self, k = 3):
        selection_ix = np.random.randint(self.n_pop)
        for ix in np.random.randint(0, len(self.pop), k-1):
            if self.scores[ix] < self.scores[selection_ix]:
                selection_ix = ix
        return self.pop[selection_ix]

    def crossover(self, p1, p2, r_cross):
        c1, c2 = p1.copy(), p2.copy()
        if np.random.rand() < r_cross:
            pt = np.random.randint(1, len(p1) - 2)
            #perform crossover
            c1 = p1[:pt] + p2[pt:]
            c2 = p2[:pt] + p1[pt:]
        return [c1, c2]
    def mutation(self, sol):
        for i in range(len(sol)):
            if np.random.rand() < self.r_mut:
                if sol[i] == 1: 
                    sol[i] = 0
                elif sol[i] == 0: 
                    sol[i] = 1
        return sol
    def gentic_algorithm(self):
        for gen in range(self.n_iter):
            #mark learning rate
            self.learnRate.append(self.best_eval) 

            #print(self.best,self.best_eval)
            #evaluate all candiates in the popultion
            self.scores = [self.objFunc(c, self.table, self.W) for c in self.pop]
            #check for new best solution
            for i in range(len(self.pop)):
                if self.scores[i] > self.best_eval:
                    self.best, self.best_eval = self.pop[i], self.scores[i]
                    print(">", gen, "new best f\n", self.pop[i],"=", self.scores[i])
            #select parents
            selected = [self.selection() for i in range(len(self.pop))]
            #create the next generation
            children = list()
            for i in range(0, len(self.pop) - 1, 2):
                #get selected paretns in pairs
                p1, p2 = selected[i], selected[i+1]
                #crossover and mutation
                for c in self.crossover(p1, p2, self.r_cross):
                    #mutation
                    self.mutation(c)
                    #store for next generation
                    children.append(c)
            #replace population
            self.pop = children

            #test
            '''
            print("-"*30, gen,"-"*30)
            print(self.pop)
            print("scores:",self.scores)
            '''
        self.getIndexObj()
        return [self.best, self.best_eval]
    def getLearnRate(self):
        rate = np.array(self.learnRate)
        plt.plot(rate)
        plt.show()
#[cost, value]
'''
table = np.array([
    [5,1],
    [4,2],
    [3,3],
    [4,2],
    [5,1]])
'''
table = np.array([[np.random.randint(0,10), np.random.randint(0,10)] for i in range(10)])
np.savetxt("table.txt", table, delimiter = ',', fmt = "%d")
#table = np.loadtxt("table.txt", int,delimiter = ',')
print(table)
#t = list([1,0,0,1,0])
#print(objFunc(t,table,10))
W = 10
h = GA(solFunc, objFunc,table,W, n_bits = len(table), n_iter = 1000, n_pop = 100, r_cross = 0.3, r_mut = 0.3)
#print(table)
sol = h.gentic_algorithm()
print("solution:")
print(sol)
print("index:")
print(h.index)
h.getLearnRate()

